import numpy as np 

class DataLoaderMultiProcess:
    def __init__(self, dataset, batch_size=2):
        pass
        self.dataset = dataset
        self.batch_size = batch_size
        self.keys = [i for i in range(len(dataset))]
        self.thread_id = 0

    def determine_shapes(self):
        # load one case
        item = self.dataset.__getitem__(0)

        num_sample = 1
        if isinstance(item, list):
            num_sample = len(item)
            item = self.dataset.__getitem__(0)[0]

        data = item["data"]
        seg = item["seg"]

        num_color_channels = data.shape[0]

        patch_size = data.shape[1:]
        data_shape = (self.batch_size * num_sample, num_color_channels, patch_size[0], patch_size[1], patch_size[2])
        seg_shape = (self.batch_size * num_sample, seg.shape[0], patch_size[0], patch_size[1], patch_size[2])
        return data_shape, seg_shape
    
    def generate_train_batch(self):
        selected_keys = np.random.choice(self.keys, self.batch_size, True, None)
        self.data_shape, self.seg_shape = self.determine_shapes()

        data_all = np.zeros(self.data_shape, dtype=np.float32)
        seg_all = np.zeros(self.seg_shape, dtype=np.int16)

        index = 0
        for j, key in enumerate(selected_keys):
            item = self.dataset.__getitem__(key)
            if isinstance(item, list):
                for i in range(len(item)):
                    data_all[index] = item[i]["data"]
                    seg_all[index] = item[i]["seg"]
                    index += 1
            else :
                data_all[index] = item["data"]
                seg_all[index] = item["seg"]
                index += 1

        return {'data': data_all, 'seg': seg_all, 'keys': selected_keys}

    def __next__(self):
        return self.generate_train_batch()

    def set_thread_id(self, thread_id):
        self.thread_id = thread_id